Deep Learning Infrastructure

LLM训练中的计算分析

计算类型、硬件分工与优化策略

核心结论
一、LLM训练中的计算类型与占比

LLM训练主要包含以下几类计算操作：

矩阵乘法 (GEMM) ~70-80%
激活函数 ~3-5%
LayerNorm ~4-6%
Softmax ~3-5%
优化器更新 ~3-5%
AllReduce通信 ~5-15%

除法操作及处理方式

Softmax: exp(x)/Σexp LayerNorm: 1/σ Attention: 1/√d Adam: 1/(√v+ε)

处理方式：GPU没有专用除法单元，除法通过倒数近似 + 乘法实现。 使用快速倒数指令（如 __frcp_rn）计算 1/x，再与被除数相乘。 这些操作由 CUDA Core 完成，属于逐元素操作。

二、硬件分工与计算占比
Tensor Core ~75%
CUDA ~17%
通信
Tensor Core (GEMM)
CUDA Core (逐元素)
通信 (AllReduce)
Tensor Core
矩阵乘法：前向/反向线性层、Attention的QK^T和·V计算
占比 ~70-80%（前向~25-30%，反向~45-50%）
CUDA Core
逐元素操作：激活函数、LayerNorm、Softmax、除法、优化器更新
占比 ~15-25%
通信
AllReduce梯度同步（CUDA Core规约 + NVLink/InfiniBand传输）
占比 ~5-15%（视并行规模而定）
核心洞察：LLM训练是矩阵乘法密集型任务，Tensor Core承担绝大部分计算量。CUDA Core负责的逐元素操作（含除法）计算量小但访存频繁，通常通过算子融合优化以减少显存读写。
01

硬件分工总览

操作类型 执行单元 原因
矩阵乘法 (GEMM) Tensor Core 专为矩阵运算优化
激活函数 (GeLU/SiLU) CUDA Core 逐元素操作
梯度计算 (GEMM部分) Tensor Core 本质是矩阵乘法
梯度计算 (逐元素) CUDA Core 如激活函数导数
除法操作 CUDA Core 倒数近似 + 乘法实现
权重更新 (Adam等) CUDA Core 逐元素操作（含除法）
AllReduce通信 CUDA Core + NVLink 数据搬运+规约
LayerNorm / Softmax CUDA Core 规约+逐元素+除法
02

前向传播详解

● 线性层 (占比最大 ~60-70%)

矩阵乘法由Tensor Core完成，偏置加法由CUDA Core完成

Y = XW + b

● 激活函数 (GeLU/SiLU)

纯逐元素操作，涉及tanh、乘法等

GeLU(x) = x · Φ(x)

● Attention计算

QK^T 和 ·V 由Tensor Core完成
Softmax/Scaling(含除法)由CUDA Core完成

softmax(QK^T/√d)·V

● LayerNorm

规约求均值/方差 + 逐元素归一化（含1/σ除法）

(x - μ) / σ · γ + β
03

反向传播详解

反向传播的梯度计算本质上也是矩阵乘法，因此主要由Tensor Core完成。 

# 对于线性层 Y = XW 的反向传播

# 权重梯度 - Tensor Core (GEMM)
dL/dW = XT · dL/dY

# 输入梯度 - Tensor Core (GEMM)
dL/dX = dL/dY · WT

# 激活函数梯度 - CUDA Core (逐元素)
dL/dx = dL/dy · f'(x)
04

权重更新 (Adam优化器)

优化器更新全部由CUDA Core完成，因为都是纯逐元素操作，包括多次除法运算。 

# Adam优化器 - 全部由CUDA Core执行

m = β₁·m + (1-β₁)·g          # 一阶动量
v = β₂·v + (1-β₂)·g²         # 二阶动量
m̂ = m / (1-β₁t)             # 偏置校正 (除法)
v̂ = v / (1-β₂t)             # 偏置校正 (除法)
w = w - lr · m̂ / (√v̂ + ε)    # 权重更新 (除法)
05

除法操作的GPU实现

GPU没有专用的除法硬件单元，所有除法操作通过倒数近似 + 乘法实现。 

# GPU除法的底层实现原理

# 计算 a / b 的过程:
reciprocal = __frcp_rn(b)    # 快速倒数近似: 1/b
result = a * reciprocal      # 乘法得到最终结果

# 对于更高精度，使用牛顿迭代优化:
x = __frcp_rn(b)             # 初始近似
x = x * (2 - b * x)          # 牛顿迭代1次
result = a * x

快速倒数的实现原理

快速倒数主要有两种实现方式：硬件查表法（GPU实际使用）和经典的位级技巧算法。

方法一：硬件查表法（GPU实际使用）

现代GPU内置查找表 (LUT)，直接根据浮点数的指数和部分尾数查表得到近似倒数： 

输入: b = 1.mantissa × 2^exp
输出: 1/b ≈ LUT[mantissa高位] × 2^(-exp-1)

硬件用几个时钟周期完成查表+插值，精度约12-13位有效位。

方法二：Fast Inverse Square Root（经典软件算法）

著名的"魔法数字"算法（Quake III引擎），利用IEEE 754浮点格式的位级技巧： 

float fast_inv_sqrt(float x) {
    int i = *(int*)&x;           // 浮点数按整数解释
    i = 0x5f3759df - (i >> 1);   // 魔法数字！
    float y = *(float*)&i;       // 整数按浮点解释
    y = y * (1.5f - 0.5f*x*y*y); // 牛顿迭代优化
    return y;
}

原理解析：

IEEE 754 单精度浮点数：(-1)^s × 1.M × 2^(E-127)

将其按整数 I 解释：I = E × 2²³ + M

对于 y = 1/√x，取对数：log(y) = -0.5 × log(x)

整数近似：I_y ≈ 1.5 × 2²³ × 127 - 0.5 × I_x

化简得：I_y = 0x5f3759df - (I_x >> 1)

两种方法对比

方法 使用场景 精度 速度
硬件LUT GPU __frcp_rn ~12位 4-8 cycles
魔法数字 CPU / 无硬件支持 ~8位 极快
牛顿迭代 精度不够时补充 每次翻倍 +4 cycles/次

GPU的 __frcp_rn 本质就是硬件查表 + 可选迭代的封装。

除法操作出处汇总

除法操作出处 具体形式 执行单元
Softmax exp(xᵢ) / Σexp(xⱼ) CUDA Core
LayerNorm (x - μ) / σ CUDA Core
Attention Scaling QK^T / √d_k CUDA Core
Adam 偏置校正 m / (1-β^t) CUDA Core
Adam 更新 m̂ / (√v̂ + ε) CUDA Core

为什么不直接用高精度除法？

方法 延迟 (cycles) 说明
硬件除法 __fdiv_rn ~20-36 完全精确
快速倒数 __frcp_rn ~4-8 足够日常使用
倒数 + 1次牛顿迭代 ~8-12 接近完全精度

深度学习对精度要求不高（FP16/BF16训练），快速倒数的精度已经足够，因此通常直接使用 __frcp_rn 省去迭代，以获得 3-5倍 的性能提升。

06

为什么Tensor Core占主导？

H100 PCIe 算力对比

Tensor Core FP16
1,979 TFLOPS
CUDA Core FP32
51 TFLOPS

LLM的计算特点：Transformer主体是GEMM操作；参数量越大，GEMM占比越高；逐元素操作（包括除法）虽然数量多，但计算量相对较小，主要瓶颈在于显存访问。

07

实际优化要点

Tensor Core利用率

矩阵维度对齐到8/16的倍数

CUDA Core操作

算子融合 (Fused Kernel)

内存带宽

Flash Attention、激活重计算

通信开销

计算-通信重叠、梯度压缩

常用术语全称

缩写 全称 说明
FMA Fused Multiply-Add 融合乘加，单指令完成 D = A×B + C
GEMM General Matrix Multiply 通用矩阵乘法，深度学习核心操作
GeLU Gaussian Error Linear Unit 高斯误差线性单元，Transformer常用激活函数
SiLU Sigmoid Linear Unit 又称Swish，SiLU(x) = x × σ(x)

算子融合与硬件的关系

算子融合是软件层面的优化策略，Tensor Core / CUDA Core是硬件执行单元。融合后的算子内部会同时调用两种硬件单元。

完整的层次关系

┌─────────────────────────────────────────────────┐
│  算子融合 (软件层面决策)                          │
│  决定哪些操作合并成一个 Kernel                    │
├─────────────────────────────────────────────────┤
│  融合后的算子 (一个 Kernel)                       │
│  __global__ void matmul_bias_gelu(...)          │
├─────────────────────────────────────────────────┤
│  Kernel 内部的硬件调度                           │
│  ├── GEMM 部分 → Tensor Core 执行               │
│  └── Bias + GeLU 部分 → CUDA Core 执行          │
│  (中间结果在寄存器传递，不写回显存)              │
├─────────────────────────────────────────────────┤
│  硬件: Tensor Core + CUDA Core                   │
└─────────────────────────────────────────────────┘

硬件级融合 vs 软件级融合

层次 名称 示例 执行单元
硬件级 FMA 指令 D = A×B + C（单条指令） Tensor Core
软件级 算子融合 GEMM + Bias + GeLU → 一个 Kernel Tensor Core + CUDA Core

Tensor Core 的基本指令 D = A×B + C 是一种硬件级 FMA 融合；而我们常说的"算子融合"是软件级优化，将多个算子合并到一个 Kernel，让 Tensor Core 输出直接传给 CUDA Core 处理，中间结果留在寄存器/共享内存中。

算子融合代码示例

# 未融合: 3次 Kernel 调用，3次显存读写
y = linear(x)      # Kernel 1: Tensor Core GEMM，写显存
y = y + bias       # Kernel 2: CUDA Core 加法，读写显存
y = gelu(y)        # Kernel 3: CUDA Core 激活，读写显存

# 融合后: 1次 Kernel 调用，1次显存写入
y = fused_linear_bias_gelu(x, weight, bias)
# 内部流程:
# 1. Tensor Core 做 GEMM → 结果留在寄存器
# 2. CUDA Core 做 bias + gelu → 结果写回显存

融合带来的性能收益

示例：Y = GeLU(X × W + b)，矩阵大小 4096×4096

❌ 未融合 (3个算子)

MatMul: 读 64MB, 写 32MB
Add: 读 32MB, 写 32MB
GeLU: 读 32MB, 写 32MB
总计: 256MB 内存访问

✅ 融合后 (1个算子)

MatMul_Bias_GeLU:
读 64MB, 写 32MB

总计: 96MB 内存访问

内存访问减少 62%，性能提升 2-3x。这也是 Flash Attention 的核心优化思路之一。